package com.example.gateway.filter;


import com.alibaba.fastjson.JSON;
import com.example.common.entity.ExcpErrResult;
import com.example.common.enums.RequestLimitEnum;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Map;


/**
 * @author Destinycc
 */
@Slf4j
@Component
public class RequestFilter extends AbstractGatewayFilterFactory<RequestFilter.Config> {

    private final MyRedisRateLimiter defaultRateLimiter;
    private final CustomKeyResolver defaultKeyResolver;

    public RequestFilter(MyRedisRateLimiter defaultRateLimiter, CustomKeyResolver defaultKeyResolver) {
        super(Config.class);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;
    }

    public CustomKeyResolver getDefaultKeyResolver() {
        return defaultKeyResolver;
    }

    public MyRedisRateLimiter getDefaultRateLimiter() {
        return defaultRateLimiter;
    }

    @SuppressWarnings("unchecked")
    @Override
    public GatewayFilter apply(Config config) {
        CustomKeyResolver resolver = (config.keyResolver == null) ? defaultKeyResolver : config.keyResolver;
        MyRedisRateLimiter limiter = (config.rateLimiter == null) ? defaultRateLimiter : config.rateLimiter;
        return (exchange, chain) -> {
            Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
            String routeId = route.getId();
            routeId = routeId + resolver.hashCode();
            String finalRouteId = routeId;
            return resolver.resolve(exchange).flatMap(key ->
                limiter.isAllowed(finalRouteId, key).flatMap(response -> {
                    for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                        exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                    }
                    if (response.isAllowed()) {
                        return chain.filter(exchange);
                    }
                    String reason = RequestLimitEnum.getDataByBeanName(resolver.getBeanName());
                    ServerHttpResponse httpResponse = exchange.getResponse();
                    httpResponse.setStatusCode(config.getStatusCode());
                    httpResponse.getHeaders().add("Content-Type", MediaType.APPLICATION_JSON_VALUE);
                    ExcpErrResult<String> result = ExcpErrResult.<String>builder().code(-1)
                            .msg("访问受限")
                            .data(reason)
                            .build();
                    String data = result.toString();
                    DataBuffer buffer = httpResponse.bufferFactory().wrap(data.getBytes(StandardCharsets.UTF_8));
                    //todo 抛异常
                    return httpResponse.writeWith(Mono.just(buffer));
                }));
        };
    }

    @Data
    public static class Config {
        private CustomKeyResolver keyResolver;
        private MyRedisRateLimiter rateLimiter;
        private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;
    }
}

